
use std::fs;
use datafusion::prelude::*;
use arrow::*;

#[global_allocator]
static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
    println!("sql_query");
    // 带列头的数据文件
    // head is : id, seq, money, bri, level, desc, flag, memo
    let csv_path = "order_table_with_head.csv";
    println!("{:?}", csv_path);

    // register the table
    let ctx = SessionContext::new();
    // let mut ctx = ExecutionContext::new();

    let mut options = CsvReadOptions::new();
    options = options.has_header(true);

    ctx.register_csv("order", csv_path, options).await?;

    // create a plan to run a SQL query
    let sql = getSQL();
    println!("sql: {}", &sql);

    let df = ctx.sql(&sql).await?;

    // execute and print results
    df.show().await?;

    let record_batch_vec = df.collect().await?;
    let count: usize = record_batch_vec
        .iter()
        .map(|v| v.num_rows())
        .sum();
    println!("query count: {}", count);





    println!("sql_filter");
    let mut options = CsvReadOptions::new();
    options = options.has_header(true);

    let df = ctx.read_csv(csv_path, options).await?;

    let df = df.filter(col("id").lt_eq(col("seq")))?
        .aggregate(vec![col("id")], vec![min(col("seq"))])?
        .limit(None, Some(100))?;

    // execute and print results, 打印展示数据
    df.show_limit(10).await?;

    // execute the plan
    let results: Vec<_> = df.collect().await?;

    let count: usize = results
        .iter()
        .map(|v| v.num_rows())
        .sum();
    println!("sql_filter show 10, query limit count: {}", count);

    // TODO 目前 pretty 编译报错
    // 格式化打印数据，  format the results
    // let pretty_results = arrow::util::pretty::pretty_format_batches(&results)?
    //     .to_string();
    // println!("sql_filter pretty_results: {}",
    //          pretty_results.trim().lines().collect::<Vec<_>>());

    Ok(())
}


fn getSQL() -> String {
    let sql;
    {
        // head is : id, seq, money, bri, level, desc, flag, memo
        sql = "SELECT id, seq, money, memo FROM order where id > 300 order by seq";
    }
    // {
    //     sql = "select * from order limit 5";
    // }
    // {
    //     sql = "select * from order \
    // union all select * from order \
    // union all select * from order \
    // union all select * from order \
    // union all select * from order ";
    // }

    sql.to_string()
}

